from osgeo import ogr,gdal
import geopandas as gpd
from PIL import Image
import numpy as np
import glob
import cv2
import os


def draw_polygon(raster_path,vector_path,field_name,out_path):
    raster_ds = gdal.Open(raster_path, gdal.GA_Update)
    cols = raster_ds.RasterXSize
    rows = raster_ds.RasterYSize
    bands = raster_ds.RasterCount
    projection = raster_ds.GetProjection()
    geotransform = raster_ds.GetGeoTransform()
    image = np.zeros((rows, cols, bands), dtype=np.uint8)
    for i in range(bands):
        band = raster_ds.GetRasterBand(i+1)
        band_array =  band.ReadAsArray()
        image[:, :, i] = band_array
    vector_ds = ogr.Open(vector_path)
    layer = vector_ds.GetLayer()
    unique_values = set()
    for feature in layer:
        value = feature.GetField(field_name)
        unique_values.add(value)
    colors = [(0,0,255),(255,0,0),(0,255,0),(0,255,255),(255,255,0),(255,0,255)]
    for index,field_value in enumerate(unique_values):
        if field_value == 0:continue
        layer.SetAttributeFilter(f"{field_name} = '{field_value}'")
        mem_driver = gdal.GetDriverByName("MEM")
        mem_raster_ds = mem_driver.Create("", cols, rows, 1, gdal.GDT_Byte)
        mem_raster_ds.SetProjection(projection)
        mem_raster_ds.SetGeoTransform(geotransform)
        gdal.RasterizeLayer(mem_raster_ds, [1], layer, burn_values=[255])
        mem_band = mem_raster_ds.GetRasterBand(1)
        mask = mem_band.ReadAsArray() 
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(image, contours, -1, colors[index], 2)
        mem_raster_ds = None  
    cv2.imwrite(out_path,image)
    raster_ds = None
    vector_ds = None
    return


def draw_color(source_path,mask_path,out_path):
    image = cv2.imdecode(np.fromfile(source_path,dtype=np.uint8),-1)
    
    palette_image  = cv2.imdecode(np.fromfile(mask_path,dtype=np.uint8),cv2.IMREAD_UNCHANGED)
    palette_image[palette_image == 0] = 0
    # palette_image = np.arange(256, dtype=np.uint8)
    
    # b, g, r, a = cv2.split(palette_image)
    # 将颜色通道转换为带有透明度的图像
    # palette_image_rgba = cv2.merge((b, g, r, a))

    # 将唯一值颜色透明度设置为80%
    alpha = 0.8
    palette_image[:, :, 3] = alpha * palette_image[:, :, 3]
    
    result_image = cv2.addWeighted(palette_image, 1, image, 1, 0)
    
    # result_image = cv2.LUT(mask, palette)
    # cv2.imshow('Result Image', result_image)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    cv2.imwrite(out_path, result_image)
    return


def clip_image(source_path,save_dir):
    image = cv2.imdecode(np.fromfile(source_path,dtype=np.uint8),-1)
    height,width = image.shape[:2]
    window_size = (1024,512)  
    step_size = (512,256)  
    for x in range(0,width, step_size[0]):
        for y in range(0,height, step_size[1]):
            window = image[y:y+window_size[1], x:x+window_size[0]]
            out_path = os.path.join(save_dir,f"{y}_{x}.png")
            cv2.imwrite(out_path,window)
    return


def join_image(source_path,data_dir,out_path):
    re_path = os.path.join(data_dir,"*.png")
    paths = glob.glob(re_path)
    src_image = cv2.imdecode(np.fromfile(source_path,dtype=np.uint8),-1)
    src_height,src_width = src_image.shape[:2]
    mask = np.ones((src_height,src_width))
    for path in paths:
        name = os.path.basename(path)
        y,x = os.path.splitext(name)[0].split("_")
        palette_image = Image.open(path)
        index_array = list(palette_image.getdata())
        width, height = palette_image.size
        index_matrix = np.array(index_array).reshape(height, width)
        tarr = mask[int(y):int(y)+height, int(x):int(x)+width]
        arr = np.maximum(tarr,index_matrix)
        mask[int(y):int(y)+height, int(x):int(x)+width] = arr    
    cv2.imwrite(out_path,mask)
    return


def filter_point(shp_path,out_shp):
    gdf = gpd.read_file(shp_path)
    gdf = gdf.sort_values(by='time')
    for i in range(10,gdf.shape[0]-11):
        array = gdf['CLASS'][i-10:i+11].values
        array = np.delete(array, [10])
        unique_values, counts = np.unique(array, return_counts=True)
        max_index = np.argmax(counts)
        is_update = counts[max_index]>= 16
        if is_update: gdf.loc[i,'CLASS'] = unique_values[max_index] 
    gdf.to_file(out_shp)
    return


if __name__ =="__main__":
    # source_path = r"/data/fengyy/houmengfei/train20240710/images_lines/mask_1_000048_1717727818041_image_draw_2600_1950_extract_2x.png"
    # mask_path = r"/data/fengyy/guiji_Data/predict_1/pseudo_color_prediction/mask_1_000048_1717727818041_image_draw_2600_1950_extract_2x.png"
    # out_path = r"/data/fengyy/guiji_Data/predict_1/mask_1_000048_1717727818041_image_draw_2600_1950_extract_2x.shp"
    shp_path = r"/data/fengyy/guiji_Data/20240716/new_test/1717120228271/1717120228271_point.shp"
    out_shp = r"/data/fengyy/guiji_Data/20240716/new_test/1717120228271/1717120228271_point_filter.shp"
    filter_point(shp_path,out_shp)
    # source_path = r"/data/fengyy/guiji_Data/test/mask_6_000424_211224202310240036134652_image_draw_transpose_2176_650_flip_horizontal_2x.png"
    # draw_color(source_path,mask_path,out_path)
    # clip_image(source_path,save_dir)
    # data_dir = r"/data/fengyy/guiji_Data/test/res"
    # save_dir = r"/data/fengyy/guiji_Data/test/tex"
    # join_image(source_path,data_dir,save_dir)